# -*- coding:utf-8 -*-

import tensorflow as tf
import os

def load_graph(frozen_graph_filename):
    #citation: code is taken from https://blog.metaflow.fr/tensorflow-how-to-freeze-a-model-and-serve-it-with-a-python-api-d4f3596b3adc#.137byfk9k
    with tf.gfile.GFile(frozen_graph_filename, "rb") as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
    with tf.Graph().as_default() as graph:
        tf.import_graph_def(graph_def,name="")
    return graph

'''适合tf1.x版本打印网络权重参数，tf2.x版本会报错AttributeError: module 'tensorflow_core._api.v2.train' has no attribute 'import_meta_graph'''
# 参考 <https://www.cnblogs.com/monologuesmw/p/13303745.html>
def txt_save(data, output_file):
    file = open(output_file, 'a')
    for i in data:
        s = str(i) + '\n'
        file.write(s)
    file.close()

def network_param(pb_path, output_file=None):
    graph = load_graph(pb_path)
    with tf.Session(graph = graph) as sess:
        variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
        for i in variables:
            print(i)     # 打印
        txt_save(variables, output_file)  # 保存txt   二选一

if __name__ == '__main__':
    checkpoint_path = './origin.pb' 
    output_file = 'network_param.txt'
    if not os.path.exists(output_file):
        network_param(checkpoint_path, output_file)
